package com.lingjtx.common.thread.core;

import io.micrometer.core.instrument.util.NamedThreadFactory;
import org.springframework.core.Ordered;
import org.springframework.stereotype.Component;

import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.*;
import java.util.logging.Level;
import java.util.logging.Logger;
import java.util.stream.Collectors;

@Component
public class ThreadPoolManager implements Ordered {

    private static final Logger log = Logger.getLogger(ThreadPoolManager.class.getName());
    private final Map<String, ExecutorService> poolMap = new ConcurrentHashMap<>();

    public ExecutorService getOrCreate(String name) {
        return poolMap.computeIfAbsent(name, this::createDefaultThreadPool);
    }

    public ExecutorService getOrCreate(String name, ThreadPoolBuilder config) {
        return poolMap.computeIfAbsent(name, config::build);
    }

    public ExecutorService get(String name) {
        return poolMap.get(name);
    }

    public void registerExecutor(String name, ExecutorService executorService) {
        ExecutorService old = poolMap.put(name, executorService);
        if (old != null) {
            old.shutdown(); // 优雅替换
        }
    }

    private ExecutorService createDefaultThreadPool(String name) {
        return new ThreadPoolBuilder().build(name);
    }

    // ========= Scheduled 类型支持 =========

    public ScheduledExecutorService getOrCreateScheduled(String name) {
        return (ScheduledExecutorService) poolMap.computeIfAbsent(name, _ -> createScheduledExecutor(name));
    }

    public ScheduledExecutorService getScheduled(String name) {
        ExecutorService executor = poolMap.get(name);
        if (executor instanceof ScheduledExecutorService) {
            return (ScheduledExecutorService) executor;
        }
        throw new IllegalStateException("线程池 [" + name + "] 不是 ScheduledExecutorService 类型");
    }

    private ScheduledExecutorService createScheduledExecutor(String name) {
        return Executors.newSingleThreadScheduledExecutor(new NamedThreadFactory(name + "-scheduler"));
    }

    // =============  线程池方法 =============

    public void execute(String name, Runnable task) {
        getOrCreate(name).execute(wrap(task));
    }

    public Future<?> submit(String name, Runnable task) {
        return getOrCreate(name).submit(wrap(task));
    }

    public <T> Future<T> submit(String name, Callable<T> task) {
        return getOrCreate(name).submit(wrap(task));
    }

    public void shutdown(String name) {
        ExecutorService pool = poolMap.get(name);
        if (pool != null) {
            pool.shutdown();
        }
    }

    public void shutdownAll(String... names) {
        for (String name : names) {
            shutdown(name);
        }
    }

    public void shutdownNow(String name) {
        ExecutorService pool = poolMap.get(name);
        if (pool != null) {
            pool.shutdownNow();
        }
    }

    public void shutdownAll() {
        poolMap.forEach((k, v) -> v.shutdown());
    }

    private Runnable wrap(Runnable task) {
        return () -> {
            try {
                task.run();
            } catch (Throwable t) {
                log.log(Level.SEVERE, "Task exception", t);
                throw new RuntimeException("Wrapped exception", t); // 包装异常
            }
        };
    }

    private <T> Callable<T> wrap(Callable<T> task) {
        return () -> {
            try {
                return task.call();
            } catch (Throwable t) {
                log.log(Level.SEVERE, "Task exception", t);
                throw t;
            }
        };
    }

    public Map<String, ExecutorService> getAllPools() {
        return Map.copyOf(poolMap);
    }

    // 监控指标接口
    public Map<String, Object> getThreadPoolMetrics() {
        return poolMap.entrySet().stream().collect(Collectors.toMap(
                Map.Entry::getKey,
                e -> {
                    if (e.getValue() instanceof ThreadPoolExecutor executor) {
                        Map<String, Object> metrics = new HashMap<>();
                        metrics.put("corePoolSize", executor.getCorePoolSize());
                        metrics.put("maximumPoolSize", executor.getMaximumPoolSize());
                        metrics.put("activeCount", executor.getActiveCount());
                        metrics.put("queueSize", executor.getQueue().size());
                        return metrics;
                    }
                    return Collections.emptyMap();
                }
        ));
    }

    @Override
    public int getOrder() {
        return 1;
    }
}